# -*- coding: utf-8 -*-
import os
import re
import logging
import maya.cmds as mc
import pymel.core as pm
import maya.mel as mel
import maya.utils as mu
from Qt import __binding__
from Qt.QtCore import *
from Qt.QtWidgets import *
import liberConfig
from liberLibs import utils

TEXTURE_NODE_ATTR_DICT = liberConfig.TEXTURE_NODE_ATTR_DICT
COLOR_SPACE_ATTR_DICT = liberConfig.COLOR_SPACE_ATTR_DICT


def get_maya_version():
    """
    get current maya version
    :return:
    """
    maya_version = mc.about(version=1)
    return maya_version


def new_file():
    mc.file(new=1, f=1)


def plugin_loaded(plugin):
    return mc.pluginInfo(plugin, q=1, loaded=1)


def load_plugin(plugin):
    if not plugin_loaded(plugin):
        mc.loadPlugin(plugin, quiet=1)
    return True


def open_file(file_name, lnr=False):
    pm.mel.eval('saveChanges("")')
    mc.file(file_name, open=1, f=1, loadNoReferences=lnr, ignoreVersion=1)


def save_as(path):
    dir_name = os.path.dirname(path)
    if not os.path.isdir(dir_name):
        os.makedirs(dir_name)
    mc.file(rename=path)
    save_file()


def save_file():
    scene_name = mc.file(q=1, sn=1)
    if scene_name.endswith(".mb"):
        mc.file(save=1, f=1, type="mayaBinary")
    elif scene_name.endswith(".ma"):
        mc.file(save=1, f=1, type="mayaAscii")


def get_history_nodes(node):
    history_nodes = mc.listHistory(node, ac=1)
    history_nodes = list(set(history_nodes))
    return history_nodes


def get_history_texture_nodes(node):
    history_nodes = get_history_nodes(node)
    texture_nodes = list()
    for typ in TEXTURE_NODE_ATTR_DICT.keys():
        try:
            texture_node = mc.ls(history_nodes, type=typ)
            if not texture_node:
                continue
            texture_nodes.extend(texture_node)
        except:pass
    return texture_nodes


def get_texture_node_of_selected_geo():
    sg_nodes = get_sg_node_of_selected()
    return get_texture_nodes_of_sg(sg_nodes)


def get_texture_nodes_of_sg(sg_nodes):
    all_texture_nodes = list()
    for sg_node in sg_nodes:
        texture_nodes = get_history_texture_nodes(sg_node)
        all_texture_nodes.extend(texture_nodes)
    texture_nodes = list(set(all_texture_nodes))
    return texture_nodes


def get_all_texture_nodes():
    sg_nodes = mc.ls(type="shadingEngine")
    return get_texture_nodes_of_sg(sg_nodes)


def export_textures(texture_nodes, texture_dir):
    if not texture_nodes:
        return
    temp_dict = dict()
    for node in texture_nodes:
        node_type = mc.nodeType(node)
        attr = "%s.%s" % (node, TEXTURE_NODE_ATTR_DICT.get(node_type))
        color_space_attr = COLOR_SPACE_ATTR_DICT.get(node_type)
        cs_attr = "%s.%s" % (node, color_space_attr) if color_space_attr else None
        temp_dict[attr] = mc.getAttr(attr)
        if cs_attr:
            temp_dict[cs_attr] = mc.getAttr(cs_attr)
    for node in texture_nodes:
        try:
            export_texture(node, texture_dir, True)
        except Exception as e:
            print "Something wrong with export textures."
            logging.error(str(e))
    return temp_dict


def export_texture(node, tex_dir, change_file_texture_name):
    node_type = mc.nodeType(node)
    attr = "%s.%s" % (node, TEXTURE_NODE_ATTR_DICT.get(node_type))
    texture = mc.getAttr(attr)
    if node_type == "file":
        texture = mc.getAttr("%s.computedFileTextureNamePattern" % node)
    if not texture:
        return
    texture = texture.replace("\\", "/")
    if not os.path.splitdrive(texture)[0]:
        texture = "%s%s" % (mc.workspace(q=1, rootDirectory=1, fullName=1), texture)
    real_path = utils.get_texture_abs_path(texture)
    if not real_path:
        return
    for each_path in real_path:
        base_name = os.path.basename(each_path)
        new_path = "%s/%s" % (tex_dir, base_name)
        if utils.copy(each_path, new_path):
            logging.info("Copy %s >> %s" % (each_path, new_path))
    if change_file_texture_name:
        texture_base_name = os.path.basename(texture)
        new_texture_path = "%s/%s" % (tex_dir, texture_base_name)
        color_space_attr = COLOR_SPACE_ATTR_DICT.get(node_type)
        cs_attr = "%s.%s" % (node, color_space_attr) if color_space_attr else None
        cs_value = mc.getAttr(cs_attr) if cs_attr else None
        mc.setAttr(attr, new_texture_path, type="string")
        if cs_attr and cs_value:
            mc.setAttr(cs_attr, cs_value, type="string")


def post_export_textures(temp_dict):
    if temp_dict:
        for attr in temp_dict:
            if temp_dict.get(attr):
                mc.setAttr(attr, temp_dict.get(attr), type="string")


def selected():
    return mc.ls(sl=1)


def export_selected(file_path, pr_flag=False):
    """
    :param file_path:a maya file path
    :param pr_flag: if True: export still as reference, else: import
    :return:
    """
    sel = mc.ls(sl=1)
    if not sel:
        print "[LIBER] info: Nothing selected."
        return
    parent_dir = os.path.dirname(file_path)
    if not os.path.isdir(parent_dir):
        os.makedirs(parent_dir)
    maya_type = get_file_type(file_path)
    mc.file(file_path, typ=maya_type, options="v=0", force=1, es=1, pr=pr_flag)


def export_obj(start, end, path, padding=4):
    parent_dir = os.path.dirname(path)
    if not os.path.isdir(parent_dir):
        os.makedirs(parent_dir)
    for i in range(start, end+1):
        mc.currentTime(i)
        prefix, suffix = os.path.splitext(path)
        new_path = "%s.%s%s" % (prefix, str(i).zfill(padding), suffix)
        mc.file(new_path, typ="OBJexport", options="groups=0;ptgroups=0;materials=0;smoothing=1;normals=1", pr=1, es=1)


def export_gpu_cache(root, directory, filename, start_time=1, end_time=1, save_multiple_files=True):
    mc.gpuCache(root, directory=directory, fileName=filename, optimize=1, writeMaterials=1, dataFormat="ogawa",
                startTime=start_time, endTime=end_time, saveMultipleFiles=save_multiple_files)


def import_gpu_cache(gpu_name, parent_name, cache_file_name):
    """
    :param gpu_name: gpu shape node name
    :param parent_name: gpu node name
    :param cache_file_name: gpu cache file name
    :return:
    """
    gpu_node = mc.createNode("gpuCache", name=gpu_name)
    mc.setAttr("%s.cacheFileName" % gpu_node, cache_file_name, type="string")
    parent_node = mc.listRelatives(gpu_node, parent=1)[0]
    final_name = mc.rename(parent_node, parent_name)
    return final_name


def get_file_type(path):
    file_type = None
    if path.endswith(".abc"):
        file_type = "Alembic"
    elif path.endswith(".mb"):
        file_type = "mayaBinary"
    elif path.endswith(".ma"):
        file_type = "mayaAscii"
    elif path.endswith(".fbx"):
        file_type = "Fbx"
    elif path.endswith(".obj"):
        file_type = "OBJexport"
    return file_type


def maya_import(file_path):
    file_type = get_file_type(file_path)
    if file_path.endswith(".obj"):
        file_type = "OBJ"
    mc.file(file_path, i=1, type=file_type, ignoreVersion=1, ra=1,
            mergeNamespacesOnClash=0, namespace=":", options="v=0", pr=1)


def export_abc(start_frame, end_frame, tar_path, root, uv_write=True, renderable_only=True, attribute=None):
    if isinstance(root, basestring):
        root = [root]
    if isinstance(attribute, basestring):
        attribute = [attribute]
    tar_dir = os.path.dirname(tar_path)
    if not os.path.isdir(tar_dir):
        os.makedirs(tar_dir)
    load_plugin("AbcExport.mll")
    j_base_string = "-frameRange {start_frame} {end_frame} -worldSpace" \
                    " -writeVisibility -file {tar_path}"
    if uv_write:
        j_base_string += " -uvWrite"
    if renderable_only:
        j_base_string += " -renderableOnly"
    if attribute:
        for attr in attribute:
            j_base_string += " -u %s" % attr
    for r in root:
        j_base_string += " -root %s" % r
    j_string = j_base_string.format(start_frame=start_frame, end_frame=end_frame, tar_path=tar_path)
    mc.AbcExport(j=j_string)


def import_abc(abc_path):
    mc.AbcImport(abc_path)


def create_reference(path, namespace_name=":", allow_repeat=False, get_group=False):
    result = None
    path = path.replace("\\", "/")
    if path.endswith(".abc"):
        load_plugin("AbcImport.mll")
    file_type = get_file_type(path)
    if allow_repeat:
        result = pm.system.createReference(path, loadReferenceDepth="all",
                                           mergeNamespacesOnClash=False,
                                           namespace=namespace_name,
                                           type=file_type)
    else:
        references = pm.listReferences()
        if not references:
            print "*" * 100
            result = pm.system.createReference(path, loadReferenceDepth="all",
                                               ignoreVersion=1, gl=1,
                                               options="v=0",
                                               mergeNamespacesOnClash=True,
                                               namespace=namespace_name,
                                               type=file_type)
        else:
            reference_paths = [ref.path for ref in references]
            if path not in reference_paths:
                result = pm.system.createReference(path, loadReferenceDepth="all",
                                                   mergeNamespacesOnClash=True,
                                                   namespace=namespace_name,
                                                   type=file_type)
            else:
                ref_node = pm.referenceQuery(path, referenceNode=1)
                if not pm.referenceQuery(ref_node, isLoaded=1):
                    pm.system.loadReference(path)
    if get_group:
        return pm.referenceQuery(result.refNode, dagPath=1, nodes=1)[0]


def export_ass(file_name, start, end):
    dir_name = os.path.dirname(file_name)
    if not os.path.isdir(dir_name):
        os.makedirs(dir_name)
    mc.file(file_name, f=1,
            options="-shadowLinks 1;-endFrame %s;-mask 2303;-lightLinks 1;-frameStep 1.0;-boundingBox;-startFrame %s" % (
            end, start),
            typ="ASS Export", pr=1, es=1)


def import_ass(dso_name):
    node = mc.createNode("aiStandIn", name="temp")
    parent_node = mc.listRelatives(node, parent=1)[0]
    parent_node = mc.rename(parent_node, "aiStandIn")
    node = mc.rename(node, "%sShape" % parent_node)
    mc.setAttr("%s.dso" % node, dso_name, type="string")
    return node


def create_ass_proxy(ass_files):
    file_num = len(ass_files)
    if file_num == 1:
        dso_name = ass_files[0]
    else:
        frame_padding = utils.frame_padding(ass_files[0], "ass")
        dso_name = re.sub("\d+\.ass$", "%s.ass" % ("#"*frame_padding), ass_files[0])
    node = import_ass(dso_name)
    if file_num > 1:
        mc.expression(s="{0}.frameNumber=(frame-1)%{1}+1".format(node, file_num))


def export_rs(file_name, start, end):
    dir_name = os.path.dirname(file_name)
    if not os.path.isdir(dir_name):
        os.makedirs(dir_name)
    mc.file(file_name, pr=1, es=1, force=1, typ="Redshift Proxy",
            options="startFrame=%s;endFrame=%s;frameStep=1;exportConnectivity=0;" % (start, end))


def import_rs(file_name):
    rs_proxy_node = mc.createNode("RedshiftProxyMesh", name="redshiftProxy")
    rs_proxy_placeholder_shape = mc.createNode("mesh", name="temp")
    transform = mc.listRelatives(rs_proxy_placeholder_shape, parent=1)[0]
    transform = mc.rename(transform, os.path.basename(file_name).split(".")[0])
    rs_proxy_placeholder_shape = mc.rename(rs_proxy_placeholder_shape, "%sShape" % transform)
    mc.connectAttr("%s.outMesh" % rs_proxy_node, "%s.inMesh" % rs_proxy_placeholder_shape, f=1)
    mc.setAttr("%s.fn" % rs_proxy_node, file_name, type="string")
    if mc.objExists("initialShadingGroup"):
        mc.sets(rs_proxy_placeholder_shape, fe="initialShadingGroup")
    return rs_proxy_node


def create_rs_proxy(rs_files):
    file_num = len(rs_files)
    if file_num == 1:
        file_name = rs_files[0]
    else:
        frame_padding = utils.frame_padding(rs_files[0], "rs")
        file_name = re.sub("\d+\.rs$", "%s.rs" % ("#"*frame_padding), rs_files[0])
    node = import_rs(file_name)
    if file_num > 1:
        mc.setAttr("%s.useFrameExtension" % node, True)
        exp = mc.createNode("expression", name="expression")
        mc.connectAttr("%s.out[0]" % exp, "%s.frameExtension" % node, f=1)
        mc.expression(exp, e=1, s="{0}.frameExtension=(frame-1)%{1}+1".format(node, file_num), ae=1, uc="all")


def get_sg_node_of_selected():
    sg_nodes = list()
    sels = selected()
    for sel in sels:
        if mc.nodeType(sel) in ["transform"]:
            shape = mc.listRelatives(sel, ad=1, c=1, fullPath=1, type="mesh")
        else:
            shape = sel
        sg_node = mc.listConnections(shape, s=0, d=1, type="shadingEngine")
        if not sg_node:
            continue
        sg_nodes.extend(sg_node)
    sg_nodes = list(set(sg_nodes))
    return sg_nodes


def create_aistandard_volume():
    shader = mc.shadingNode("aiStandardVolume", asShader=1, asRendering=1)
    sg = mc.shadingNode("shadingEngine", asRendering=1, name="%sSG" % shader)
    mc.connectAttr("%s.out" % shader, "%s.vs" % sg)
    return sg


def maya_version():
    return mc.about(version=1)


def export_xgen_archive(name, directory, start, end):
    import xgenm
    file_name = mc.file(q=1, sn=1)
    sels = selected()
    hlp = xgenm.xmaya.xgmArchiveExport.xgmArchiveExport()
    version = maya_version()
    if version in ["2015"]:
        hlp.processDir(name, directory, [file_name], 0, '0.0', '0.0', start, end)
    else:
        hlp.processDir(name, directory, [file_name], sels, 0, '0.0', '0.0', start, end)


def import_xgen_archive(path):
    import xgenm.ui as xgui
    xgen_description_editor = xgui.createDescriptionEditor(showIt=False)
    primitive_tab = xgen_description_editor.primitiveTab.widget()
    archive_files_ui = primitive_tab.files
    file_array = [path]
    addArchives(archive_files_ui, file_array)


def addArchives(archive_files_node, files):
    '''Add multiple archives from arc files.'''
    import xgenm
    groups = archive_files_node.archiveGroups()
    importedGroups = []
    for f in files:
        with open(f) as fp:
            t = fp.read()
            lines = t.splitlines()
            group = archive_files_node.archiveGroupsFromLines(lines)
            importedGroups.extend(group)
            groups.extend(group)

    if xgenm.xgGlobal.Maya:
        materialFiles = []
        de = xgenm.xgGlobal.DescriptionEditor
        curPal = de.currentPalette()
        curDesc = de.currentDescription()
        for g in importedGroups:
            resolvedMaterials = xgenm.XgExternalAPI.findFileInXgDataPath(str(g.materials), "", curPal, curDesc,
                                                                         "ArchivePrimitive", False, 0, 0, 0)
            if resolvedMaterials and len(resolvedMaterials):
                materialFiles.append(resolvedMaterials)
        if len(materialFiles):
            strFiles = ""
            for m in materialFiles:
                strFiles += m + "\n"
            for f in materialFiles:
                sceneName = os.path.basename(f)
                sceneExt = os.path.splitext(sceneName)
                if len(sceneExt[1]):
                    sceneName = sceneExt[0]
                mc.file(f, i=True, type="mayaAscii", ignoreVersion=True, ra=True, mergeNamespacesOnClash=False,
                        namespace=sceneName, options="v=0;p=17;f=0", pr=True)
    archive_files_node.setArchiveGroups(groups)


def import_arnold_vdb(file_name):
    volume_node = mc.createNode("aiVolume", name="temp")
    parent_node = mc.listRelatives(volume_node, p=1)[0]
    parent_node = mc.rename(parent_node, "aiVolume")
    volume_node = mc.rename(volume_node, "%sShape" % parent_node)
    mc.setAttr("%s.filename" % volume_node, file_name, type="string")
    mc.setAttr("%s.v" % volume_node, k=0)
    mc.setAttr("%s.covm[0]" % volume_node, 0, 1, 1)
    mc.setAttr("%s.cdvm[0]" % volume_node, 0, 1, 1)
    mc.setAttr("%s.grids" % volume_node, "density", type="string")
    return volume_node


def import_rs_vdb(file_name):
    volume_node = mc.createNode("RedshiftVolumeShape", name="temp")
    parent_node = mc.listRelatives(volume_node, p=1)[0]
    parent_node = mc.rename(parent_node, "rsVolume")
    volume_node = mc.rename(volume_node, "%sShape" % parent_node)
    mc.setAttr("%s.fn" % volume_node, file_name, type="string")
    mc.setAttr("%s.v" % volume_node, k=0)
    mc.setAttr("%s.covm[0]" % volume_node, 0, 1, 1)
    mc.setAttr("%s.cdvm[0]" % volume_node, 0, 1, 1)
    mc.setAttr("%s.vir" % volume_node, 1)
    mc.setAttr("%s.vif" % volume_node, 1)
    return volume_node


def create_rs_volume():
    shader = mc.shadingNode("RedshiftVolume", asShader=1, asRendering=1)
    mc.setAttr("%s.density_name" % shader, "density", type="string")
    sg = mc.shadingNode("shadingEngine", asRendering=1, name="%sSG" % shader)
    mc.connectAttr("%s.oc" % shader, "%s.vs" % sg)
    return sg


def delete_unused_nodes():
    mel.eval("MLdeleteUnused;")


def get_maya_win(module="PySide"):
    """
    get a QMainWindow Object of maya main window
    :param module (optional): string "PySide"(default) or "PyQt4"
    :return main_window: QWidget or QMainWindow object
    """
    import maya.OpenMayaUI as mui
    prt = mui.MQtUtil.mainWindow()
    if module == "PyQt":
        import sip
        main_window = sip.wrapinstance(long(prt), QObject)
    elif module in ["PySide"]:
        if __binding__ in ["PySide"]:
            import shiboken
        elif __binding__ in ["PySide2"]:
            import shiboken2 as shiboken
        main_window = shiboken.wrapInstance(long(prt), QWidget)
    elif module == "mayaUI":
        main_window = "MayaWindow"
    else:
        raise ValueError('param "module" must be "mayaUI" "PyQt4" or "PySide"')
    return main_window


def get_current_renderer():
    renderer = mc.getAttr('defaultRenderGlobals.currentRenderer')
    return renderer


def show_as_panel(widget_instance, title=None):
    if not isinstance(widget_instance, QWidget):
        raise ValueError("%s is not a Qt Widget." % widget_instance)

    obj_name = widget_instance.objectName()
    # parent current dialog to maya win
    maya_main_window = get_maya_win()
    widget_instance.setParent(maya_main_window)
    # dock the panel on maya
    if not title:
        title = widget_instance.windowTitle() or widget_instance.objectName()
    dock_panel(obj_name, widget_instance, title)

    return widget_instance


def dock_panel(object_name, widget_instance, title):
    maya_panel_name = "panel_%s" % object_name
    # delete existed panel
    if mc.control(maya_panel_name, query=True, exists=True):
        mc.deleteUI(maya_panel_name)

    if int(str(mc.about(api=True))[:4]) < 2017:

        # Create a new Maya window.
        maya_window = mc.window()

        # Add a layout to the Maya window.
        maya_layout = mc.formLayout(parent=maya_window)

        # Reparent the Shotgun app panel widget under the Maya window layout.
        mc.control(object_name, edit=True, parent=maya_layout)

        # Keep the Shotgun app panel widget sides aligned with the Maya window layout sides.
        mc.formLayout(maya_layout,
                      edit=True,
                      attachForm=[(object_name, 'top', 1),
                                  (object_name, 'left', 1),
                                  (object_name, 'bottom', 1),
                                  (object_name, 'right', 1)])

        # Dock the Maya window into a new tab of Maya Channel Box dock area.
        mc.dockControl(maya_panel_name, area="right", content=maya_window, label=title)

        # Once Maya will have completed its UI update and be idle,
        # raise (with "r=True") the new dock tab to the top.
        mu.executeDeferred("cmds.dockControl('%s', edit=True, r=True)" % maya_panel_name)

    else:  # Maya 2017 and later
        # Delete any default workspace control state that might have been automatically
        # created by Maya when a previously existing Maya panel was closed and deleted.
        if mc.workspaceControlState(maya_panel_name, exists=True):
            mc.workspaceControlState(maya_panel_name, remove=True)

        # Retrieve the Channel Box dock area, with error reporting turned off.
        # This MEL function is declared in Maya startup script file UIComponents.mel.
        # It returns an empty string when this dock area cannot be found in the active Maya workspace.
        dock_area = mel.eval('getUIComponentDockControl("Channel Box / Layer Editor", false)')

        # This UI script will be called to build the UI of the new dock tab.
        # It will embed the Shotgun app panel widget into a Maya workspace control.
        # Maya 2017 expects this script to be passed in as a string, not as a function pointer.
        # See function _build_workspace_control_ui() below for a commented version of this script.
        ui_script = "import pymel.core as pm\n" \
                    "workspace_control = pm.setParent(query=True)\n" \
                    "pm.control('%s', edit=True, parent=workspace_control)" \
                    % object_name

        # The following UI script can be used for development and debugging purposes.
        # This script has to retrieve and import the current source file in order to call
        # function _build_workspace_control_ui() below to build the workspace control UI.
        # ui_script = "import imp\n" \
        #             "panel_generation = imp.load_source('%s', '%s')\n" \
        #             "panel_generation._build_workspace_control_ui('%s')" \
        #             % (__name__, __file__.replace(".pyc", ".py"), object_name)

        # Give an initial width to the docked Shotgun app panel widget when first shown.
        # Otherwise, the workspace control would use the width of the currently displayed tab.
        size_hint = widget_instance.sizeHint()
        if size_hint.isValid():
            # Use the widget layout preferred size.
            widget_width = size_hint.width()
        else:
            # Since no size is recommended for the widget, use its current width.
            widget_width = widget_instance.width()

        # Dock the Shotgun app panel widget into a new tab of the Channel Box dock area.
        # When this dock area was not found in the active Maya workspace,
        # the Shotgun app panel widget is embedded into a floating workspace control window.
        # This floating workspace control can then be docked into an existing dock area by the user.
        dock_tab = mc.workspaceControl(maya_panel_name,
                                       tabToControl=(dock_area, -1),  # -1 to append a new tab
                                       uiScript=ui_script,
                                       loadImmediately=True,
                                       retain=False,  # delete the dock tab when it is closed
                                       label=title,
                                       initialWidth=widget_width,
                                       minimumWidth=True,  # set the minimum width to the initial width
                                       r=True  # raise the new dock tab to the top
                                       )

        # Now that the workspace dock tab has been created, let's update its UI script.
        # This updated script will be saved automatically with the workspace control state
        # in the Maya layout preference file when the user will choose to quit Maya,
        # and will be executed automatically when Maya is restarted later by the user.

        # The script will delete the empty workspace dock tab that Maya will recreate on startup
        # when the user previously chose to quit Maya while the panel was opened.
        deferred_script = "import maya.cmds as cmds\\n" \
                          "if cmds.workspaceControl('%(id)s', exists=True):\\n" \
                          "    cmds.deleteUI('%(id)s')" \
                          % {"id": maya_panel_name}

        # The previous script will need to be executed once Maya has completed its UI update and be idle.
        ui_script = "import maya.utils\n" \
                    "maya.utils.executeDeferred(\"%s\")\n" \
                    % deferred_script

        # Update the workspace dock tab UI script.
        mc.workspaceControl(maya_panel_name, edit=True, uiScript=ui_script)


if __name__ == "__main__":
    pass
